#!/home/gop/zelixir/Apps/miniconda3/bin/python
import sys, os
from Bio.PDB import Select, PDBIO
from Bio.PDB.PDBParser import PDBParser
import argparse
import  urllib.request

class ChainSelect(Select):
    def __init__(self, chain):
        self.chain = chain

    def accept_chain(self, chain):
        if chain.get_id() == self.chain:
            return 1
        else:
            return 0

def mkdir(path):
    folder = os.path.exists(path)
    if not folder:
        os.makedirs(path)

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Prepare input file', usage='use "%(prog)s --help" for more information')
    parser.add_argument('-i', '--input',  default='XXX', help='[required] pdb file')
    parser.add_argument('-p', '--pdb_id', default='XXX', help='[required] pdb file')
    parser.add_argument('-c', '--chain',  default='A',   help='the selected chain id')
    parser.add_argument('-o', '--output', default='./',  help='path to output the created pdb file')

    args = parser.parse_args()

    chain_id = args.chain
    outpath  = args.output
    mkdir(outpath)

    if args.input == 'XXX' and args.pdb_id != 'XXX':
        pdb_id = args.pdb_id.lower()
        pdb_base = pdb_id
        f = urllib.request.urlopen("http://www.rcsb.org/pdb/files/"+pdb_id+".pdb")
        pdb_in  = '{}/{}.pdb'.format(outpath, pdb_id)
        pdb_out = '{}/{}_{}.pdb'.format(outpath, pdb_id, chain_id)
        with open(pdb_in,"wb") as g:
            g.write(f.read())
    elif args.input != 'XXX' and args.pdb_id == 'XXX':
        pdb_in   = args.input
        pdb_base = os.path.basename(pdb_in)
        pdb_out  = '{}/{}_{}.pdb'.format(outpath, pdb_base[:-4], chain_id)
    else:
        print('no legal input')
        exit(1)

    p = PDBParser(PERMISSIVE=1)
    structure = p.get_structure(pdb_base, pdb_in)

    chains = []
    for c in structure.get_chains():
        chains.append(c.get_id())

    if chain_id in chains:
        io_w_no_h = PDBIO()
        io_w_no_h.set_structure(structure)
        io_w_no_h.save(pdb_out, ChainSelect(chain_id))
    else:
        print('The choice of chain should be one of the following:')
        print(chains)
        exit(1)

